package com.ace.lab.starter.launcher;

import com.ace.lab.base.util.NetTools;
import com.ace.lab.starter.config.RpcServerProperties;
import com.ace.lab.starter.constant.ZookeeperPath;
import com.ace.lab.starter.exception.RpcServerBootException;
import com.ace.lab.starter.exception.ServiceRegistryException;
import com.ace.lab.starter.meta.ExposedServiceMetaData;
import com.ace.lab.starter.utils.RpcClassHelper;
import io.grpc.BindableService;
import io.grpc.Server;
import io.grpc.netty.NettyServerBuilder;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import org.apache.curator.framework.CuratorFramework;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.aop.support.AopUtils;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.boot.context.event.ApplicationReadyEvent;
import org.springframework.context.ApplicationListener;

import java.io.IOException;
import java.net.Inet4Address;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;

public class RpcServerLauncher implements ApplicationListener<ApplicationReadyEvent>, BeanPostProcessor, DisposableBean {
    private final Logger logger = LoggerFactory.getLogger(RpcServerLauncher.class);

    private RpcServerProperties producerProperties;
    private CuratorFramework curatorFramework;
    private Server server;
    private List<ExposedServiceMetaData> serviceMetaDataList = new ArrayList<>();
    @Value("${org.springframework.boot.test.context.SpringBootTestContextBootstrapper:false}")
    private boolean testContextBootStrapper;

    public RpcServerLauncher(RpcServerProperties producerProperties, CuratorFramework curatorFramework) {
        this.producerProperties = producerProperties;
        this.curatorFramework = curatorFramework;
    }

    @Override
    public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
        return bean;
    }

    @Override
    @SuppressWarnings("unchecked")
    public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
        if (!testContextBootStrapper) {
            Optional<Class<?>> serviceImplBaseClass = getServiceImplBaseClass(bean);
            serviceImplBaseClass.ifPresent(clazz -> genServiceMetaData(bean, (Class<? extends BindableService>) clazz));
        }
        return bean;
    }

    private Optional<Class<?>> getServiceImplBaseClass(Object bean) {
        Class<?> serviceImpClass;
        if (AopUtils.isAopProxy(bean)) {
            serviceImpClass = AopUtils.getTargetClass(bean);
        } else {
            serviceImpClass = bean.getClass();
        }
        if (bean instanceof BindableService && serviceImpClass.getSuperclass().getName().endsWith("ImplBase")) {
            return Optional.of(serviceImpClass.getSuperclass());
        }
        return Optional.empty();
    }

    @Override
    public void onApplicationEvent(ApplicationReadyEvent event) {
        if (!testContextBootStrapper) {
            launchRpcServer();
            registryService();
        }
    }

    private void launchRpcServer() {
        SocketAddress address = new InetSocketAddress(NetTools.getLocalAddress(addr -> addr instanceof Inet4Address), producerProperties.getPort());
        NettyServerBuilder builder = NettyServerBuilder.forAddress(address)
                .channelType(NioServerSocketChannel.class)
                .bossEventLoopGroup(new NioEventLoopGroup(producerProperties.getAcceptThreadCount()))
                .workerEventLoopGroup(new NioEventLoopGroup(producerProperties.getNioThreadCount()));
        for (ExposedServiceMetaData metaData : serviceMetaDataList) {
            builder.addService((BindableService) metaData.getBean());
        }
        server = builder.build();
        try {
            server.start();
        } catch (IOException e) {
            throw new RpcServerBootException("launch rpc server error.", e);
        }
        logger.info("grpc server running on port: {}", producerProperties.getPort());
        Runtime.getRuntime().addShutdownHook(new Thread(() -> {
            logger.info("*** shutting down gRPC server since JVM is shutting down");
            RpcServerLauncher.this.destroy();
            logger.info("*** gRPC server shut down");
        }));
    }

    private void registryService() {
        logger.info("service registry start...");
        String address = NetTools.getLocalAddress(addr -> addr instanceof Inet4Address).getHostAddress() + ":" + producerProperties.getPort();
        try {
            String path;
            String serviceName;
            for (ExposedServiceMetaData serviceMetaData : serviceMetaDataList) {
                serviceName = serviceMetaData.getName();
                path = ZookeeperPath.SERVICE_ROOT_PATH + serviceName;
                Stat stat = curatorFramework.checkExists().forPath(path);
                if (stat == null) {
                    curatorFramework.create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT).forPath(path);
                }
                curatorFramework.create().withMode(CreateMode.EPHEMERAL).forPath(path + "/" + address);
                logger.info("registed service: {}", serviceName);
            }
        } catch (Exception e) {
            logger.error("registry service error:{}", e);
            throw new ServiceRegistryException("registry service error", e);
        }
        logger.info("service registry done...");
    }

    private void genServiceMetaData(Object bean, Class<? extends BindableService> serviceImplBaseClass) {
        String serviceName = RpcClassHelper.genServiceName(serviceImplBaseClass);
        ExposedServiceMetaData metaData = new ExposedServiceMetaData();
        metaData.setName(serviceName);
        metaData.setBean(bean);
        serviceMetaDataList.add(metaData);
    }

    @Override
    public void destroy() {
        if (server != null && !server.isShutdown()) {
            server.shutdown();
        }
    }

}
